from tkinter.filedialog import askdirectory
from tkinter import scrolledtext
from xml.dom.minidom import parse
import tkinter as tk
import tkinter.ttk
import tkinter
import threading
from osgeo import gdal
from osgeo import gdalconst
import zipfile
import datetime
import shutil
import copy
import os
import re


# 获取1c含云量
def read_cloud(zipFilePath):
	""""""
	xmlFilePathRe = "^S2.*\.SAFE/MTD_MSIL1C\.xml$"
	zipFile = zipfile.ZipFile(zipFilePath, "r")
	zipFileInfos = zipFile.infolist()
	for zipFileInfo in zipFileInfos:
		if re.match(xmlFilePathRe, zipFileInfo.filename):
			xmlFilePath = zipFile.open(zipFileInfo.filename, "r")
			domTree = parse(xmlFilePath)
			collection = domTree.documentElement
			cloudPercentElement = collection.getElementsByTagName("Cloud_Coverage_Assessment")[0]
			cloudPercent = round(float(copy.deepcopy(cloudPercentElement.childNodes[0].data)))
	return cloudPercent

# unzip zip to file
def unzip_file(zip_path, nameres, save_dir):
	""""""
	paths = []
	zipFile = zipfile.ZipFile(zip_path, "r")
	for namere in nameres:
		for info in zipFile.infolist():
			if re.match(namere, os.path.basename(info.filename)):
				path = zipFile.extract(info.filename, save_dir)
				paths.append(path)
	zipFile.close()
	return paths

# raster to attribute
def get_raster_attribute(raster_path):
	""""""
	dataset = gdal.Open(raster_path)
	attributes = {
		"RasterXSize": dataset.__getattr__("RasterXSize"),
		"RasterYSize": dataset.__getattr__("RasterYSize"),
		"RasterCount": dataset.__getattr__("RasterCount"),
		"SpatialRef": dataset.GetSpatialRef(),
		"GeoTransform": dataset.GetGeoTransform(),
		"DataType": dataset.GetRasterBand(1).DataType,
		"NoData": dataset.GetRasterBand(1).GetNoDataValue()}
	return attributes

# raster to array
def get_raster_band_array(raster_path, shape):
	""""""
	dataset = gdal.Open(raster_path, 1)
	(buf_xsize, buf_ysize) = shape
	xsize = dataset.RasterXSize
	ysize = dataset.RasterYSize
	array = dataset.ReadAsArray(
		xoff=0,
		yoff=0,
		xsize=xsize,
		ysize=ysize,
		buf_xsize=buf_xsize,
		buf_ysize=buf_ysize,
		resample_alg=gdalconst.GRIORA_NearestNeighbour,
		interleave='band')
	return array

# data to raster file
def get_raster_file(raster_path, arrays, attribute):
	""""""
	driver = gdal.GetDriverByName("GTiff")
	dataset = driver.Create(
		raster_path,
		xsize=attribute["RasterXSize"],
		ysize=attribute["RasterYSize"],
		bands=attribute["RasterCount"],
		eType=attribute["DataType"]
	)
	dataset.SetSpatialRef(attribute["SpatialRef"])
	dataset.SetGeoTransform(attribute["GeoTransform"])
	nodata = 0 if attribute["NoData"] is None else attribute["NoData"]
	for index in range(1, attribute["RasterCount"] + 1):
		band = dataset.GetRasterBand(index)
		band.WriteArray(arrays[index-1])
		band.SetNoDataValue(nodata)
	return

# single band rasters to mutil band raster
def get_mutil_band_raster(source_paths, target_path):
	""""""
	count = len(source_paths)

	attributes = []
	for source_path in source_paths:
		attribute = get_raster_attribute(source_path)
		attributes.append(attribute)

	# 多线程的读取栅格文件属性
	# with ThreadPoolExecutor(max_workers=count) as executor:
	# 	attributes = executor.map(get_raster_attribute, source_paths)

	# 获得模板属性
	attributes = list(attributes)
	pixelsizes = [attribute["GeoTransform"][1] for attribute in attributes]
	attribute = attributes[pixelsizes.index(min(pixelsizes))]
	shape = (attribute["RasterXSize"], attribute["RasterYSize"])
	attribute.update({"RasterCount": count})

	# 多线程的读取栅格文件数组
	# shapes = [shape] * count
	# with ThreadPoolExecutor(max_workers=5) as executor:
	# 	arrays = executor.map(get_raster_band_array, source_paths, shapes)

	arrays = [
		get_raster_band_array(source_path, shape)
		for source_path in source_paths
	]

	# 数据写入栅格
	arrays = list(arrays)
	get_raster_file(target_path, arrays, attribute)
	return

# get raster file
def get_muitl_band_sentinel_raster(args):
	[zip_path, tif_path] = args
	bands = [2,3,4,5,6,7,8,12]
	zip_name = os.path.basename(zip_path)
	string = f"{datetime.datetime.now().strftime('%H:%M:%S')}>success {zip_name}\n"
	if os.path.exists(tif_path):
		return string
	nameres = [
		f"T.*B{'0'*(len(f'{index}') < 2)}{index}\.jp2"
		for index in bands
	]
	save_dir = os.path.dirname(tif_path)
	safe_name = zip_name.replace(".zip",".SAFE")
	safe_dir = os.path.join(save_dir, safe_name)
	if os.path.exists(safe_dir):
		shutil.rmtree(safe_dir)
	band_paths = unzip_file(zip_path, nameres, save_dir)

	get_mutil_band_raster(band_paths, tif_path)	
	
	if os.path.exists(safe_dir):
		shutil.rmtree(safe_dir)
	string = f"{datetime.datetime.now().strftime('%H:%M:%S')}>success {zip_name}\n"
	return string


# ui
class my_ui:

	width, height = 850, 700
	line_height = 60

	width, height = 850, 600
	thread_num = 10


	def __init__(self):
		self.get_window()
		self.zip_folder = tk.StringVar()

		self.tif_folder = tk.StringVar()
		self.thread_num = tk.IntVar()
		# self.thread_num.set(1)

		line_iter = iter(range(50))
		self.zip_folder_frame(next(line_iter))
		self.tif_folder_frame(next(line_iter))
		self.start_frame(next(line_iter))
		self.progressbar_frame(next(line_iter))
		self.print_frame(next(line_iter))
		self.window.mainloop()

	def select_zip_folder(self):
		_path = askdirectory()
		_path = _path.replace('/','\\')
		self.zip_folder.set(_path)
		return

	def select_tif_folder(self):
		_path = askdirectory()
		_path = _path.replace('/','\\')
		self.tif_folder.set(_path)
		return

	def get_window(self):
		self.window = tk.Tk()
		self.window.title('ZIP转TIF数据工具')
		screenwidth = self.window.winfo_screenwidth() # 屏幕宽度
		screenheight = self.window.winfo_screenheight() # 屏幕高度
		x = int((screenwidth - self.width) / 2)
		y = int((screenheight - self.height) / 2)
		locat = f"{self.width}x{self.height}+{x}+{y}"
		self.window.geometry(locat) # 大小以及位置
		self.window.resizable(0,0) 
		return

	def zip_folder_frame(self, row):
		frame=tkinter.ttk.Frame(self.window,width=self.width)
		frame.place(x=0,y=self.line_height*row)
		tk.Label(
            frame,text="ZIP目录",font=('宋体','14','bold'),height=3,
        ).grid(column=0,row=0,sticky='we')  #定义一个标签lable
		self.zip_dir_txt=tk.Entry(
			frame,
			width=60,
			background = 'white',
            borderwidth = 3,
            relief = 'sunken',
            insertwidth = 3,
            font = ('', '14', 'bold'),
            textvariable=self.zip_folder,
            selectforeground = 'red',
            selectbackground = 'blue',
            selectborderwidth = 0,
            show=None
			)
		self.zip_dir_txt.grid(column=1,row=0,sticky='we')	
		self.select_zip_dir_button=tk.Button(
			frame,
			text="选择目录",
			font=('','14','bold'), 
			bg="#50b7c1",
			command=self.select_zip_folder
		) 
		self.select_zip_dir_button.grid(column=2,row=0,sticky='we')
		return

	def tif_folder_frame(self, row):
		frame=tkinter.ttk.Frame(self.window,width=self.width)
		frame.place(x=0,y=self.line_height*row)
		tk.Label(
            frame,text="TIF目录",font=('宋体','14','bold'),height=3,
        ).grid(column=0,row=0,sticky='we')  #定义一个标签lable
		self.tif_dir_txt=tk.Entry(
			frame,
			width=60,
			background = 'white',
            borderwidth = 3,
            relief = 'sunken',
            insertwidth = 3,
            font = ('', '14', 'bold'),
            textvariable=self.tif_folder,
            selectforeground = 'red',
            selectbackground = 'blue',
            selectborderwidth = 0,
            show=None
		)
		self.tif_dir_txt.grid(column=1,row=0,sticky='we')     
		tk.Button(
			frame,
			text="选择目录",
			font=('','14','bold'), 
			bg="#50b7c1",
			command=self.select_tif_folder
		).grid(column=2,row=0,sticky='we')
		return

	def start_frame(self, row):
		frame=tkinter.ttk.Frame(self.window)
		frame.place(x=0,y=self.line_height*row-10)
		self.start_button=tk.Button(
			frame,
			text="开始合成",
			font=('宋体','14','bold'),	
			bg="#50b7c1",
			command=lambda: self.thread_it(self.run_main),
		)
		self.start_button.pack(padx=self.width/2-10,pady=20)
		return

	def progressbar_frame(self,row):
		frame=tkinter.ttk.Frame(self.window)
		frame.place(x=0,y=self.line_height*row)
		tk.Label(
            frame,text="运行进度",font=('宋体','14','bold'),height=3,
        ).grid(column=0,row=0,sticky='we')  
		self.progressbar = tkinter.ttk.Progressbar(
			frame,
			length=740,
			value=0,
            cursor='spider',
            mode="determinate",
            # mode="indeterminate",
			orient=tk.HORIZONTAL,
		)
		self.progressbar.grid(column=1,row=0,sticky='we')
		return

	def print_frame(self, row):
		frame=tkinter.ttk.Frame(self.window,width=self.width)
		frame.place(x=0,y=self.line_height*row)
		tk.Label(
            frame,text="运行信息",font=('宋体','14','bold'),
        ).grid(column=0,row=0,sticky='we')
		self.scrolled_text = scrolledtext.ScrolledText(
			frame, 
			state='normal',
			wrap = tk.WORD,
			font=('Times New Roman', 12),
			width=90, 
			height=18,
		) 
		self.scrolled_text.grid(column=1,row=0,rowspan=100,sticky='we')
		return

	def thread_it(self,func,*args):
		# 线程
		t = threading.Thread(target=func, args=args)
		# 守护 !!!
		t.setDaemon(True)
		# 启动
		t.start()
		# 阻塞--卡死界面！
		# t.join()
		return

	def run_main(self):

		zip_dir = self.zip_dir_txt.get()
		if not os.path.exists(zip_dir):
			return

		tif_dir = self.tif_dir_txt.get()
		if not os.path.exists(tif_dir):
			return

		zip_paths = [
			os.path.join(zip_dir,name)
			for name in os.listdir(zip_dir)
			if re.match("^S2[AB]_MSIL1C_.*\.zip$",name)
		]

		string = f"{datetime.datetime.now().strftime('%H:%M:%S')}>开始合成！\n"
		self.scrolled_text.insert(tk.END, string)
		self.scrolled_text.see(tk.END)
		self.window.update()

		self.progressbar["maximum"] = len(zip_paths) 
		for index,zip_path in enumerate(zip_paths):
			zip_name = os.path.basename(zip_path)
			tile = zip_name.split("_")[5][1:]
			date = zip_name.split("_")[2][:8]
			code = zip_name.split("_")[6][9:-4]
			cloud = read_cloud(zip_path)
			tif_name = f"L1C_{tile}_{date}-{code}_{cloud}.tif"
			tif_path = os.path.join(tif_dir, tif_name)
			string = get_muitl_band_sentinel_raster([zip_path, tif_path])
			self.scrolled_text.insert(tk.END, string)
			self.scrolled_text.see(tk.END)
			self.progressbar["value"] = index+1
			self.window.update()

		tk.messagebox.showinfo('Message', '任务完成！')
		return


if __name__ == "__main__":
	my_ui()


